package com.havey.ye.controller;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.RequestResponseAdvisor;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import com.havey.ye.dto.ChatDTO;
import com.havey.ye.entity.ChatRecords;
import com.havey.ye.service.IChatRecordsService;

import jakarta.annotation.Resource;
import reactor.core.publisher.Flux;

@RestController
@RequestMapping("/aigc")
public class ChatController {
    /*
     * 定义log
     */
    private static final Logger log = LoggerFactory.getLogger(ChatController.class);

    @Resource
    ChatClient chatClient;
    @Resource
    ChatModel chatModel;
    @Resource
    IChatRecordsService icChatRecordsService;
    /**
     * 1.根据sessionId查询是否有数据，如果有数据，从全局map中获得chatClient;
     */
    private final Map<String, ChatClient> chatClientMap = new ConcurrentHashMap<>();

    /**
     * 1.根据sessionId查询是否有数据，如果有数据，从全局map中获得chatClient;
     * 2. 如果没有数据，设置一个chatClient，并保存到全局map中。
     * 2.1 设置chatClient的方法
     * ChatMemory chatMemory = new InMemoryChatMemory();
     * chatClient = ChatClient.builder(chatModel)
     * .defaultAdvisors(new MessageChatMemoryAdvisor(chatMemory))
     * .build();
     * 
     * @param chatDTO
     * @return
     */
    @PostMapping("/chat")
    public Flux<String> chat(@RequestBody ChatDTO chatDTO) {
        String sessionId = chatDTO.getSessionId();
        String message = chatDTO.getQuestionContent();
        String userName = chatDTO.getUserName();
        log.info("Chat request from user: {}", userName);
        ChatClient chatClient = chatClientMap.get(sessionId);
        if (chatClient == null) {
            ChatMemory chatMemory = new InMemoryChatMemory();
            List<RequestResponseAdvisor> requestResponseAdvisors = new ArrayList<>();
            requestResponseAdvisors.add(new MessageChatMemoryAdvisor(chatMemory));
            // requestResponseAdvisors.add(new SimplelogAdvisor());
            chatClient = ChatClient.builder(chatModel)
                    .defaultAdvisors(requestResponseAdvisors)
                    .build();
            chatClientMap.put(sessionId, chatClient);
        }
        StringBuilder sb = new StringBuilder();
        Flux<String> result = chatClient.prompt()
                .user(message)
                .stream()
                .content()
                .doOnNext(response -> {
                    log.info("Model response: {}", response);
                    sb.append(response);
                })
                .doOnComplete(() -> {
                    ChatRecords chatRecords = new ChatRecords();
                    chatRecords.setSessionId(sessionId);
                    chatRecords.setQuestionContent(message);
                    chatRecords.setAnswerContent(sb.toString()); // Set answer content
                    chatRecords.setCreatedBy(userName);
                    icChatRecordsService.save(chatRecords);
                });
        return result;
    }

    /**
     * 1. 入参是一个sessionId.
     * 2. 根据sessionId，把chatClientMap删除掉
     */
    @PostMapping("/endSession")
    public void endSession(String sessionId) {
        chatClientMap.remove(sessionId);
    }
}